from PyQt5 import QtCore, QtWidgets,QtGui
import cv2
import time
class DrawingBoard(QtWidgets.QLabel):
    def __init__(self, parent = None):
        super().__init__(parent)
        self.src_img = None
        self.show_pix = None
        self.drawing_rect = [0,0,0,0]
        self.painter = QtGui.QPainter()
        self.mouse_down = False
        self.old_w = None
        self.old_h = None
        
    def clearup(self):
        self.drawing_rect = [0,0,0,0]
        self.drawing()

    def getROI(self):
        px = self.old_w / self.src_img.size().width()
        py = self.old_h / self.src_img.size().height()

        return [int(px*self.drawing_rect[0] + 0.5), 
        int(py*self.drawing_rect[1]+0.5),
        int( px*self.drawing_rect[2]+0.5),
        int( py*self.drawing_rect[3]+0.5)]

    def drawing(self):
        if self.src_img is None:
            return

        self.painter.begin(self.show_pix)
        self.painter.drawImage(0, 0, self.src_img)
        drawing_rect = self.drawing_rect
        if drawing_rect[0] + drawing_rect[1] + drawing_rect[2] + drawing_rect[3] != 0:
            self.painter.setPen(QtGui.QColor(255, 0, 0))
            self.painter.drawRect(QtCore.QRect(QtCore.QPoint(drawing_rect[0], drawing_rect[1]), QtCore.QPoint(drawing_rect[2], drawing_rect[3])))
        self.painter.end()
        self.setPixmap(self.show_pix)

    def setImage(self, result, is_gray):
        self.old_w = result.shape[1]
        self.old_h = result.shape[0]

        if is_gray == False:
            result = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
            qImg = QtGui.QImage(result.data, 
                                result.shape[1], 
                                result.shape[0],
                                result.shape[1]*3,
                                QtGui.QImage.Format_RGB888)
            if qImg.size().width() > self.size().width() or qImg.size().height() > self.size().height():
                qImg = qImg.scaled(self.width(),self.height(), QtCore.Qt.KeepAspectRatio)
            self.show_pix = QtGui.QPixmap.fromImage(qImg)
            self.src_img = qImg
            time.sleep(0.1)
            self.drawing()
        else:
            qImg = QtGui.QImage(result.data, 
                                result.shape[1], 
                                result.shape[0],
                                result.shape[1],
                                QtGui.QImage.Format_Grayscale8)
            if qImg.size().width() > self.size().width() or qImg.size().height() > self.size().height():
                qImg = qImg.scaled(self.width(),self.height(), QtCore.Qt.KeepAspectRatio)
            self.show_pix = QtGui.QPixmap.fromImage(qImg)
            self.src_img = qImg
            
            self.drawing()

    def mouseMoveEvent(self, event):
        if self.mouse_down:
            x = event.pos().x()
            if x < 0:
                x = 0
            elif x > self.src_img.size().width():
                x = self.src_img.size().width() - 2

            y = event.pos().y()
            if y < 0:
                y = 0
            elif y > self.src_img.size().height():
                y = self.src_img.size().height() - 2
            self.drawing_rect[2] = x
            self.drawing_rect[3] = y
            self.drawing()
        

    def mousePressEvent(self, event):
        self.mouse_down = True
        self.drawing_rect[0] = event.pos().x()
        self.drawing_rect[1] = event.pos().y()
        self.drawing_rect[2] = event.pos().x()
        self.drawing_rect[3] = event.pos().y()

    def mouseReleaseEvent(self, event):
        self.mouse_down = False
        x = event.pos().x()
        if x < 0:
            x = 0
        elif x > self.src_img.size().width():
            x = self.src_img.size().width() - 2

        y = event.pos().y()
        if y < 0:
            y = 0
        elif y > self.src_img.size().height():
            y = self.src_img.size().height() - 2
        self.drawing_rect[2] = x
        self.drawing_rect[3] = y
        self.drawing()